package org.usmile.algorithms.learn.graph.weighted;

import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;
import java.util.PriorityQueue;

public class Prim {
    public List<WeightedEdge> prim(WeightedAdjSet graph) {
        boolean[] visited = new boolean[graph.getV()];
        visited[0] = true;

        List<WeightedEdge> result = new ArrayList<>();
        // 进行V-1次切分，每次切分得到最短的横切边
        for (int i = 0; i < graph.getV(); i++) {
            for (int v = 0; v < graph.getV(); v++) {
                if (visited[v]) {
                    int minWeight = Integer.MAX_VALUE;
                    int minW = -1;
                    for (int w : graph.adj(v)) {
                        // 找到最小横切边
                        if (!visited[w] && graph.getWeight(v, w) < minWeight) {
                            minW = w;
                            minWeight = graph.getWeight(v, w);
                        }
                    }
                    if (minW != -1) {
                        result.add(new WeightedEdge(v, minW, minWeight));
                        visited[minW] = true;
                    }
                }
            }
        }

        return result;
    }

    public List<WeightedEdge> prim1(WeightedAdjSet graph) {
        boolean[] visited = new boolean[graph.getV()];
        visited[0] = true;
        List<WeightedEdge> result = new ArrayList<>();
        PriorityQueue<WeightedEdge> priorityQueue = new PriorityQueue<>(Comparator.comparingInt(WeightedEdge::getWeight));
        for (int w : graph.adj(0)) {
            priorityQueue.add(new WeightedEdge(0, w, graph.getWeight(0, w)));
        }

        while (!priorityQueue.isEmpty()) {
            WeightedEdge minEdge = priorityQueue.poll();
            // 不是横切边
            if (visited[minEdge.getV()] && visited[minEdge.getW()]) {
                continue;
            }
            result.add(minEdge);
            // 扩展切分
            int newV = visited[minEdge.getV()] ? minEdge.getW() : minEdge.getV();
            visited[newV] = true;
            for (int w : graph.adj(newV)) {
                if (!visited[w]) {
                    priorityQueue.add(new WeightedEdge(newV, w, graph.getWeight(newV, w)));
                }
            }
        }

        return result;
    }

    public static void main(String[] args) {
        WeightedAdjSet graph = new WeightedAdjSet(4);
        graph.addEdge(0, 1, 1);
        graph.addEdge(0, 3, 2);
        graph.addEdge(1, 2, 2);
        graph.addEdge(1, 3, 4);
        graph.addEdge(2, 3, 3);

        System.out.println(new Prim().prim1(graph));
    }
}
